from firewall_interface import *
import IPy

#@author:  fhl cjh

# ====================生成命令======================
#-----要求-----
#- 根据已有sample-data，创建Rule对象，判断数据有效性，生成iptables命令
#- 封装函数 
#       def addRule(rule:Rule)  -- 根据-I ；判断数据是否有效，拼接字符串
#       def deleteRule(chain:Direction, num:int)  

#-验证 使用subprocess验证是否能够添加和删除rule

    # 判断数据是否有效 src port， -1   0.0.0.0/0
    # num是否在正确范围内；
    # 是否需要添入命令

#Rule对象的使用  
#字符串的操作
#正则表达式
#format格式化输出

#----常用命令----
    #在链中第num条规则之前插入规则，默认是1
        #iptables -I [INPUT/OUTPUT][num] -j[target] -s [src] --sport [] -d [dst] --dport []

    #删除链中第num条命令
        #iptables -D [INPUT/OUTPUT] [num]

    #列出某条链中的规则
        #iptables -nvL [INPUT/OUTPUT] --line-numbers


#前端会发过来的数据
sample_data_1 = {
  'direction':0,          #INPUT/OUTPUT 
  'target': 1 ,           #Target -j
  'num':1,                #-A -I  0最高，直接-I
  'protocol': 6,          #-p
  'sport': -1,            #--sport
  'dport': 8000,          #--dport 
  'src': "0.0.0.0/0",     #-s
  'dst': "0.0.0.0/0",     #-d
}

sample_data_2 = {
  'direction':1 ,
  'target': 0 ,
  'num':1,
  'protocol': 1 ,
  'sport': -1,
  'dport': -1,
  'src': "192.168.100.114",
  'dst': "0.0.0.0/0",
}


#要删除的rule
sample_data_3 = {
  'num' : 1,
  'chain' : Direction.INPUT
}

#判断ip地址，端口等的有效性，并生成iptables命令
def getAddRuleCmd(newRule:Rule):

  direction = newRule.direction
  num = newRule.num
  target= newRule.target
  protocol = newRule.protocol
  src = newRule.src 
  dst =  newRule.dst
  dport = newRule.dport
  sport = newRule.sport

  try:
    if src != None:
      IPy.IP(src)
    if dst != None:
      IPy.IP(dst)

  except ValueError:
    return -1
  
  if direction != Direction.OUTPUT and direction != Direction.INPUT :
    return -1
    
  if target != Target.ACCEPT and target != Target.REJECT :
    return -1
    
  if protocol != Protocol.ICMP_V4 and protocol != Protocol.TCP and protocol != Protocol.UDP :
    return -1


  cmd = ''
  if direction == 1 : 
    direction='-I OUTPUT' 
  else : 
    direction='-I INPUT'

  if target == 1 : 
    target='-j REJECT'
  else : 
    target='-j ACCEPT'

  if protocol==Protocol.ICMP_V4 : protocol ='-p icmp'
  elif protocol==Protocol.TCP   : protocol ='-p tcp'
  elif protocol==Protocol.UDP   : protocol ='-p udp'

  if src.startswith('0.0.0.0') or src == None: 
    src =''
  else:
    src = "-s {}".format(src) 

  if dst.startswith('0.0.0.0') or dst == None: 
    dst =''
  else:
    dst = "-d {}".format(dst) 

  if(dport == -1): 
    dport=''
  else: 
    dport = "--dport {}".format(dport)

  if(sport == -1): 
    sport=''
  else: 
    sport = "--sport {}".format(sport)

  if protocol == 'icmp':
    cmd = 'iptables {} {} {} {} {} {}'.format(direction,num,target,protocol,src,dst)
  else:
    cmd = "iptables {} {} {} {} {} {} {} {}".format(direction, num, target, protocol, src, dst, dport, sport)
  
  #     iptables -I [INPUT/OUTPUT][num] -j[target] -s [src] --sport [] -d [dst] --dport []

  print(cmd)

  return cmd

def getDeleteRuleCmd (chain:Direction, num:int):
  if chain == 0 or chain == 1:
    if chain == 1 : 
      chain ='OUTPUT'
    else : 
      chain='INPUT'
    cmd = 'iptables -D {} {}'.format(chain,num)
    return cmd


#test
# print( getDeleteRuleCmd(sample_data_3['chain'],sample_data_3['num']) )
# print( getAddRuleCmd(Rule(sample_data_2)) )
# print( getAddRuleCmd(Rule(sample_data_1)) )

